import torch
from torch.utils.data import TensorDataset
from torch.utils.data import DataLoader
import torch.nn as nn
from torchsummary import summary
import torch.optim as optim
from sklearn.model_selection import train_test_split
import numpy as np
import pandas as pd
import time

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
def create_dataset():
    data = pd.read_csv('./手机价格预测.csv')
    x,y = data.iloc[:,:-1],data.iloc[:,-1]
    x = x.astype(np.float32)
    x_train,x_valid,y_train,y_valid = train_test_split(x,y,test_size=0.2,random_state=42)
    train_dataset = TensorDataset(torch.tensor(x_train.values),torch.tensor(y_train.values))
    valid_dataset = TensorDataset(torch.tensor(x_valid.values),torch.tensor(y_valid.values))
    return train_dataset,valid_dataset,x_train.shape[1],len(y.unique())

class MobileNet(nn.Module):
    def __init__(self,input_dim,output_dim):
        super().__init__()
        self.linear1 = nn.Linear(input_dim,128)
        self.linear2 = nn.Linear(128,256)
        self.out = nn.Linear(256,output_dim)

    def forward(self,x):
        x = torch.relu(self.linear1(x))
        x = torch.relu(self.linear2(x))
        output = self.out(x)
        return output

def train():
    train_dataset,valid_dataset,input_size,num_classes = create_dataset()
    train_dataloader = DataLoader(train_dataset,batch_size=8,shuffle=True,drop_last= True)
    model = MobileNet(input_size,num_classes).to( device= device)
    model.train()
    optimizer = optim.Adam(model.parameters(),lr=1e-3)
    criterion = nn.CrossEntropyLoss()
    epochs = 50
    loss_list = []
    for epoch in range(epochs):
        start_time = time.time()
        total = time.time()
        total_loss = 0.0
        total_iter = 0
        for train_x,train_y in train_dataloader:
            train_x,train_y = train_x.to(device),train_y.to(device)
            output = model(train_x)
            loss = criterion(output,train_y)
            total_loss += loss.item()
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            total_iter += 1
        print('epoch:',epoch,'loss:',total_loss/total_iter,'time:',time.time()-start_time)
    torch.save(model.state_dict(),'model/PhonemobileNet.pth')
def inference():
    train_dataset,valid_dataset,input_size,num_classes = create_dataset()
    test_dataloader = DataLoader(valid_dataset,batch_size=8,shuffle=True)
    model = MobileNet(input_size,num_classes).to(device=device)
    model.load_state_dict(torch.load('model/PhonemobileNet.pth',map_location=lambda storage,loc: storage))
    model.eval()
    correct = 0
    for test_x,test_y in test_dataloader:
        test_x,test_y = test_x.to(device),test_y.to(device)
        with torch.no_grad():
            output = model(test_x)
        prd_idex = torch.argmax(output,dim=-1)
        correct += torch.sum(prd_idex == test_y)
    print('acc:',correct.item()/len(valid_dataset))

if __name__ == '__main__':

    train()
    inference()
